"""
黄金交易预测项目 - 机器学习模型模块
包含多种机器学习算法和模型评估
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, TimeSeriesSplit, cross_val_score
from sklearn.preprocessing import StandardScaler, RobustScaler
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier, GradientBoostingRegressor, GradientBoostingClassifier
from sklearn.linear_model import LinearRegression, LogisticRegression, Ridge, Lasso
from sklearn.svm import SVR, SVC
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, accuracy_score, precision_score, recall_score, f1_score, classification_report, confusion_matrix
import xgboost as xgb
import lightgbm as lgb
import joblib
import os
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')


class GoldMLModels:
    """黄金交易预测机器学习模型类"""
    
    def __init__(self, model_dir="models"):
        self.model_dir = model_dir
        os.makedirs(model_dir, exist_ok=True)
        self.scaler = None
        self.models = {}
        self.feature_importance = {}
        
    def prepare_data(self, data, target_col='target_return_1d', test_size=0.2, scale_features=True):
        """
        准备训练和测试数据
        
        Args:
            data (pd.DataFrame): 特征数据
            target_col (str): 目标列名
            test_size (float): 测试集比例
            scale_features (bool): 是否标准化特征
            
        Returns:
            tuple: (X_train, X_test, y_train, y_test, feature_cols)
        """
        print("准备数据...")
        
        # 移除包含NaN的行
        df = data.copy()
        
        # 选择特征列（排除目标变量和非特征列）
        exclude_cols = ['date', 'open', 'high', 'low', 'close', 'volume', 'gold_close', 'usd_close']
        exclude_cols += [col for col in df.columns if 'target_' in col]
        
        feature_cols = [col for col in df.columns if col not in exclude_cols]
        
        # 移除包含NaN的特征列
        valid_features = []
        for col in feature_cols:
            if df[col].isna().sum() / len(df) < 0.5:  # 如果NaN比例小于50%
                valid_features.append(col)
        
        feature_cols = valid_features
        print(f"选择了 {len(feature_cols)} 个特征")
        
        # 准备特征和目标
        X = df[feature_cols].copy()
        y = df[target_col].copy()
        
        # 移除包含NaN的行
        valid_idx = ~(X.isna().any(axis=1) | y.isna())
        X = X[valid_idx]
        y = y[valid_idx]
        
        print(f"有效数据样本数: {len(X)}")
        
        # 时间序列分割（按时间顺序）
        split_idx = int(len(X) * (1 - test_size))
        X_train = X.iloc[:split_idx]
        X_test = X.iloc[split_idx:]
        y_train = y.iloc[:split_idx]
        y_test = y.iloc[split_idx:]
        
        print(f"训练集大小: {len(X_train)}")
        print(f"测试集大小: {len(X_test)}")
        
        # 特征标准化
        if scale_features:
            print("标准化特征...")
            self.scaler = RobustScaler()
            X_train_scaled = self.scaler.fit_transform(X_train)
            X_test_scaled = self.scaler.transform(X_test)
            
            X_train = pd.DataFrame(X_train_scaled, columns=X_train.columns, index=X_train.index)
            X_test = pd.DataFrame(X_test_scaled, columns=X_test.columns, index=X_test.index)
        
        return X_train, X_test, y_train, y_test, feature_cols
    
    def train_regression_models(self, X_train, y_train, model_names=['rf', 'xgb', 'lgb', 'ridge']):
        """
        训练回归模型
        
        Args:
            X_train (pd.DataFrame): 训练特征
            y_train (pd.Series): 训练目标
            model_names (list): 要训练的模型列表
            
        Returns:
            dict: 训练好的模型字典
        """
        print("训练回归模型...")
        
        models = {}
        
        for model_name in model_names:
            print(f"训练 {model_name} 模型...")
            
            if model_name == 'rf':
                model = RandomForestRegressor(
                    n_estimators=100,
                    max_depth=10,
                    min_samples_split=5,
                    min_samples_leaf=2,
                    random_state=42,
                    n_jobs=-1
                )
            elif model_name == 'xgb':
                model = xgb.XGBRegressor(
                    n_estimators=100,
                    max_depth=6,
                    learning_rate=0.1,
                    subsample=0.8,
                    colsample_bytree=0.8,
                    random_state=42,
                    n_jobs=-1
                )
            elif model_name == 'lgb':
                model = lgb.LGBMRegressor(
                    n_estimators=100,
                    max_depth=6,
                    learning_rate=0.1,
                    subsample=0.8,
                    colsample_bytree=0.8,
                    random_state=42,
                    n_jobs=-1,
                    verbose=-1
                )
            elif model_name == 'ridge':
                model = Ridge(alpha=1.0, random_state=42)
            elif model_name == 'lasso':
                model = Lasso(alpha=1.0, random_state=42)
            elif model_name == 'linear':
                model = LinearRegression()
            elif model_name == 'svr':
                model = SVR(kernel='rbf', C=1.0, gamma='scale')
            else:
                continue
            
            # 训练模型
            model.fit(X_train, y_train)
            models[model_name] = model
            
            # 保存模型
            model_path = os.path.join(self.model_dir, f"{model_name}_regression.joblib")
            joblib.dump(model, model_path)
            
            print(f"{model_name} 模型训练完成并保存到 {model_path}")
        
        self.models.update(models)
        return models
    
    def train_classification_models(self, X_train, y_train, model_names=['rf', 'xgb', 'lgb', 'logistic']):
        """
        训练分类模型
        
        Args:
            X_train (pd.DataFrame): 训练特征
            y_train (pd.Series): 训练目标
            model_names (list): 要训练的模型列表
            
        Returns:
            dict: 训练好的模型字典
        """
        print("训练分类模型...")
        
        models = {}
        
        for model_name in model_names:
            print(f"训练 {model_name} 分类模型...")
            
            if model_name == 'rf':
                model = RandomForestClassifier(
                    n_estimators=100,
                    max_depth=10,
                    min_samples_split=5,
                    min_samples_leaf=2,
                    random_state=42,
                    n_jobs=-1
                )
            elif model_name == 'xgb':
                model = xgb.XGBClassifier(
                    n_estimators=100,
                    max_depth=6,
                    learning_rate=0.1,
                    subsample=0.8,
                    colsample_bytree=0.8,
                    random_state=42,
                    n_jobs=-1
                )
            elif model_name == 'lgb':
                model = lgb.LGBMClassifier(
                    n_estimators=100,
                    max_depth=6,
                    learning_rate=0.1,
                    subsample=0.8,
                    colsample_bytree=0.8,
                    random_state=42,
                    n_jobs=-1,
                    verbose=-1
                )
            elif model_name == 'logistic':
                model = LogisticRegression(random_state=42, max_iter=1000)
            elif model_name == 'svc':
                model = SVC(kernel='rbf', C=1.0, gamma='scale', random_state=42)
            else:
                continue
            
            # 训练模型
            model.fit(X_train, y_train)
            models[model_name] = model
            
            # 保存模型
            model_path = os.path.join(self.model_dir, f"{model_name}_classification.joblib")
            joblib.dump(model, model_path)
            
            print(f"{model_name} 分类模型训练完成并保存到 {model_path}")
        
        self.models.update(models)
        return models
    
    def evaluate_regression_models(self, X_test, y_test):
        """
        评估回归模型
        
        Args:
            X_test (pd.DataFrame): 测试特征
            y_test (pd.Series): 测试目标
            
        Returns:
            pd.DataFrame: 评估结果
        """
        print("评估回归模型...")
        
        results = []
        
        for model_name, model in self.models.items():
            if 'classification' in str(type(model)):
                continue
                
            try:
                # 预测
                y_pred = model.predict(X_test)
                
                # 计算指标
                mse = mean_squared_error(y_test, y_pred)
                rmse = np.sqrt(mse)
                mae = mean_absolute_error(y_test, y_pred)
                r2 = r2_score(y_test, y_pred)
                
                # 方向准确率
                direction_acc = np.mean((y_test > 0) == (y_pred > 0))
                
                results.append({
                    'Model': model_name,
                    'MSE': mse,
                    'RMSE': rmse,
                    'MAE': mae,
                    'R2': r2,
                    'Direction_Accuracy': direction_acc
                })
                
                print(f"{model_name}: RMSE={rmse:.6f}, R2={r2:.4f}, Direction_Acc={direction_acc:.4f}")
                
            except Exception as e:
                print(f"评估 {model_name} 时出错: {e}")
        
        results_df = pd.DataFrame(results)
        results_df = results_df.sort_values('RMSE')
        
        return results_df
    
    def evaluate_classification_models(self, X_test, y_test):
        """
        评估分类模型
        
        Args:
            X_test (pd.DataFrame): 测试特征
            y_test (pd.Series): 测试目标
            
        Returns:
            pd.DataFrame: 评估结果
        """
        print("评估分类模型...")
        
        results = []
        
        for model_name, model in self.models.items():
            if 'classification' not in str(type(model)) and 'Classifier' not in str(type(model)):
                continue
                
            try:
                # 预测
                y_pred = model.predict(X_test)
                y_pred_proba = None
                
                if hasattr(model, 'predict_proba'):
                    y_pred_proba = model.predict_proba(X_test)[:, 1]
                
                # 计算指标
                accuracy = accuracy_score(y_test, y_pred)
                precision = precision_score(y_test, y_pred, average='weighted', zero_division=0)
                recall = recall_score(y_test, y_pred, average='weighted', zero_division=0)
                f1 = f1_score(y_test, y_pred, average='weighted', zero_division=0)
                
                results.append({
                    'Model': model_name,
                    'Accuracy': accuracy,
                    'Precision': precision,
                    'Recall': recall,
                    'F1': f1
                })
                
                print(f"{model_name}: Accuracy={accuracy:.4f}, F1={f1:.4f}")
                
            except Exception as e:
                print(f"评估 {model_name} 时出错: {e}")
        
        results_df = pd.DataFrame(results)
        results_df = results_df.sort_values('Accuracy', ascending=False)
        
        return results_df
    
    def get_feature_importance(self, model_name, feature_cols):
        """
        获取特征重要性
        
        Args:
            model_name (str): 模型名称
            feature_cols (list): 特征列名
            
        Returns:
            pd.DataFrame: 特征重要性
        """
        if model_name not in self.models:
            print(f"模型 {model_name} 不存在")
            return pd.DataFrame()
        
        model = self.models[model_name]
        
        try:
            if hasattr(model, 'feature_importances_'):
                importance = model.feature_importances_
            elif hasattr(model, 'coef_'):
                importance = np.abs(model.coef_)
                if len(importance.shape) > 1:
                    importance = importance.mean(axis=0)
            else:
                print(f"模型 {model_name} 不支持特征重要性")
                return pd.DataFrame()
            
            feature_importance = pd.DataFrame({
                'Feature': feature_cols,
                'Importance': importance
            }).sort_values('Importance', ascending=False)
            
            self.feature_importance[model_name] = feature_importance
            return feature_importance
            
        except Exception as e:
            print(f"获取 {model_name} 特征重要性时出错: {e}")
            return pd.DataFrame()
    
    def predict(self, model_name, X):
        """
        使用指定模型进行预测
        
        Args:
            model_name (str): 模型名称
            X (pd.DataFrame): 特征数据
            
        Returns:
            np.array: 预测结果
        """
        if model_name not in self.models:
            print(f"模型 {model_name} 不存在")
            return None
        
        model = self.models[model_name]
        return model.predict(X)
    
    def save_results(self, results, filename):
        """保存评估结果"""
        filepath = os.path.join(self.model_dir, filename)
        results.to_csv(filepath, index=False, encoding='utf-8-sig')
        print(f"结果已保存到: {filepath}")
    
    def load_model(self, model_name, model_type='regression'):
        """加载保存的模型"""
        model_path = os.path.join(self.model_dir, f"{model_name}_{model_type}.joblib")
        if os.path.exists(model_path):
            model = joblib.load(model_path)
            self.models[model_name] = model
            print(f"模型 {model_name} 加载成功")
            return model
        else:
            print(f"模型文件不存在: {model_path}")
            return None


def main():
    """主函数 - 测试机器学习模型"""
    print("=" * 60)
    print("黄金交易预测机器学习模型程序")
    print("=" * 60)
    
    # 加载特征数据
    features_file = "data/features/gold_features.csv"
    
    if not os.path.exists(features_file):
        print(f"特征数据文件不存在: {features_file}")
        return
    
    print("加载特征数据...")
    features_data = pd.read_csv(features_file)
    print(f"加载数据: {features_data.shape}")
    
    # 创建模型训练器
    ml_models = GoldMLModels()
    
    # 准备回归数据
    print("\n=== 回归任务 ===")
    X_train, X_test, y_train, y_test, feature_cols = ml_models.prepare_data(
        features_data, 
        target_col='target_return_1d',
        test_size=0.2
    )
    
    # 训练回归模型
    regression_models = ml_models.train_regression_models(
        X_train, y_train, 
        model_names=['rf', 'xgb', 'lgb', 'ridge']
    )
    
    # 评估回归模型
    regression_results = ml_models.evaluate_regression_models(X_test, y_test)
    print("\n回归模型评估结果:")
    print(regression_results)
    
    # 保存回归结果
    ml_models.save_results(regression_results, 'regression_results.csv')
    
    # 获取最佳模型的特征重要性
    best_reg_model = regression_results.iloc[0]['Model']
    print(f"\n最佳回归模型: {best_reg_model}")
    feature_importance = ml_models.get_feature_importance(best_reg_model, feature_cols)
    
    if not feature_importance.empty:
        print(f"\n{best_reg_model} 模型前10个重要特征:")
        print(feature_importance.head(10))
        
        # 保存特征重要性
        importance_file = os.path.join(ml_models.model_dir, f'{best_reg_model}_feature_importance.csv')
        feature_importance.to_csv(importance_file, index=False, encoding='utf-8-sig')
    
    # 准备分类数据
    print("\n=== 分类任务 ===")
    X_train_cls, X_test_cls, y_train_cls, y_test_cls, feature_cols_cls = ml_models.prepare_data(
        features_data, 
        target_col='target_direction_1d',
        test_size=0.2
    )
    
    # 训练分类模型
    classification_models = ml_models.train_classification_models(
        X_train_cls, y_train_cls, 
        model_names=['rf', 'xgb', 'lgb', 'logistic']
    )
    
    # 评估分类模型
    classification_results = ml_models.evaluate_classification_models(X_test_cls, y_test_cls)
    print("\n分类模型评估结果:")
    print(classification_results)
    
    # 保存分类结果
    ml_models.save_results(classification_results, 'classification_results.csv')
    
    # 获取最佳分类模型的特征重要性
    best_cls_model = classification_results.iloc[0]['Model']
    print(f"\n最佳分类模型: {best_cls_model}")
    cls_feature_importance = ml_models.get_feature_importance(best_cls_model, feature_cols_cls)
    
    if not cls_feature_importance.empty:
        print(f"\n{best_cls_model} 模型前10个重要特征:")
        print(cls_feature_importance.head(10))
        
        # 保存特征重要性
        cls_importance_file = os.path.join(ml_models.model_dir, f'{best_cls_model}_feature_importance.csv')
        cls_feature_importance.to_csv(cls_importance_file, index=False, encoding='utf-8-sig')
    
    print("\n机器学习模型训练和评估完成！")
    print("=" * 60)


if __name__ == "__main__":
    main()
